# -*- coding: utf-8 -*-


"""
@File    : 06_文件转换之ASC转TIF-批量.py
@Author  : fungis@163.com
@notice  : 文件类型asc转TIF
"""

import os
import sys

import numpy as np
from osgeo import gdal


# 将asc矩阵文件转换为TIF
def gis_asc_to_tif(projection, tran, celltype, input_file, output_file):
    """
    将数组转成tif文件写入硬盘
    :param projection: 投影信息
    :param tran: 几何信息
    :param celltype: 值类型  float or int
    :param input_file: 文件输入路径
    :param output_file: 文件输出路径
    :return:
    """

    if os.path.exists(input_file) is False:
        raise Exception('[Errno 2] 该文件不存在: \'' + input_file + '\'')

    num = 6  # 读取前6行（如果是5行的数据，这里修改为5）
    metadata_list = []  # 元数据列表
    with open(input_file, 'r') as rows_input_file:
        for i in range(num):
            tem_row = str(rows_input_file.readline()).strip()  # 读取两端剔除空格的字符串
            list_str = tem_row.split(" ")
            item_tem = list_str[-1]  # 取list的最后一个元素，即相应的头文件值
            metadata_list.append(item_tem)
    # 根据读取结果取值
    ncols = int(metadata_list[0])
    nrows = int(metadata_list[1])
    xllcorner = float(metadata_list[2])
    yllcorner = float(metadata_list[3])
    cellsize = float(metadata_list[4])  # 像元大小
    NODATA_value = float(metadata_list[5])  # 空值默认值
    xltcorner = xllcorner
    yltcorner = yllcorner + cellsize * nrows
    # 转换位置信息更改为asc文件制定的脚点
    new_tran = list(tran)
    new_tran[0] = xltcorner
    new_tran[3] = yltcorner
    geo_transform = new_tran

    # 矩阵提取
    ascii_grid = np.loadtxt(input_file, skiprows=6)
    row = ascii_grid.shape[0]  # 矩阵的行数
    columns = ascii_grid.shape[1]  # 矩阵的列数

    # print(geo_transform)
    # dim_z = mat.shape[2]  # 通道数
    dim_z = 1
    driver = gdal.GetDriverByName('GTiff')  # 创建驱动
    # 创建文件
    if celltype == "float":
        dst_ds = driver.Create(output_file, columns, row, dim_z, gdal.GDT_Float32)
    elif celltype == "int":
        dst_ds = driver.Create(output_file, columns, row, dim_z, gdal.GDT_UInt16)
    else:
        raise Exception('[Errno 4] 未设置值类型或设置字符串存在错误！')
    dst_ds.SetGeoTransform(geo_transform)  # 设置几何信息
    # 坐上角坐标更新
    dst_ds.SetProjection(projection)  # 设置投影

    # 将数组的各通道写入tif图片
    map = ascii_grid[:, :]
    band = dst_ds.GetRasterBand(1)
    # band.SetNoDataValue(-9999)
    band.SetNoDataValue(NODATA_value)
    dst_ds.GetRasterBand(1).WriteArray(map)  # 单波段

    # 多波段
    # for channel in np.arange(dim_z):
    #     map = mat[:, :, channel]
    #     dst_ds.GetRasterBand(int(channel + 1)).WriteArray(map)

    dst_ds.FlushCache()  # 写入硬盘
    dst_ds = None


# 获取投影信息
def get_pro_from_img(img_file_path):
    """
    读取栅格数据，将其转换成对应数组
    :param: img_file_path: 栅格数据路径
    :return: 返回投影，几何信息，和转换后的数组(5888,5888,10)
    """
    dataset = gdal.Open(img_file_path)  # 读取栅格数据
    # print('处理图像的栅格波段数总共有：', dataset.RasterCount)
    # 判断是否读取到数据
    if dataset is None:
        raise Exception('[Errno 3] 无法读取该文件: \'' + img_file_path + '\'')
        sys.exit(1)  # 退出
    projection = dataset.GetProjection()  # 投影
    transform = dataset.GetGeoTransform()  # 几何信息
    # 直接读取dataset
    # img_array = dataset.ReadAsArray()
    return projection, transform


filter_list = ['.ASC', '.asc']  # 设置过滤后的文件类型 当然可以设置多个类型


def find_filter_file(dirname):
    result = []  # 所有的文件
    for maindir, subdir, file_name_list in os.walk(dirname):
        # print("提示：当前主目录下的所有目录", "2:", subdir)  # 当前主目录下的所有目录
        print("3:", file_name_list)  # 当前主目录下的所有文件

        for filename in file_name_list:
            apath = os.path.join(maindir, filename)  # 合并成一个完整路径
            ext = os.path.splitext(apath)[1]  # 获取文件后缀 [0]获取的是除了文件名以外的内容
            # 设置过滤器，过滤不符合格式的文件
            if ext in filter_list:
                result.append(apath)
    return result


if __name__ == '__main__':

    input_file_path = r'./data-use/asc/'  # 输入路径
    output_file_path = r'./results/tif/'  # 输出路径
    if not os.path.exists(output_file_path):
        # os.makedirs可以创建路径上多个
        os.makedirs(output_file_path)
    file_list = find_filter_file(input_file_path)
    if len(file_list) > 0:
        for raster_path in file_list:
            # 投影信息
            projection = 'GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563,AUTHORITY["EPSG","7030"]],AUTHORITY["EPSG","6326"]],PRIMEM["Greenwich",0,AUTHORITY["EPSG","8901"]],UNIT["degree",0.0174532925199433,AUTHORITY["EPSG","9122"]],AXIS["Latitude",NORTH],AXIS["Longitude",EAST],AUTHORITY["EPSG","4326"]]'
            # 坐标信息
            tran = (-180.0, 0.5, 0.0, 90.0, 0.0, -0.5)
            # 输出文件路径
            name = os.path.basename(raster_path).split('.')[0]
            result_file = output_file_path + name + '.tif'  # 输出文件位置
            del name  # 删除无效变量
            gis_asc_to_tif(projection, tran, 'float', raster_path, result_file)
        print('------处理完成------')
    else:
        raise Exception('[Errno 1] 该路径下无{}类型文件'.format(filter_list))
